#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Author  : hu_cl
# @Contact : 760730895@qq.com
# @Date    : 2020/10/23 15:38
# @File    : process_thread_vulnerability.py
import datetime
import json
import random
import os

import config
import pymysql
import requests
import time
from math import floor
import multiprocessing
import sched

s = sched.scheduler(time.time, time.sleep)
req_url_base = 'https://redqueen.tj-un.com/'  # 漏洞主网址
leak_url = '/IntelDetails.html?id='
req_header_vlist = {
    # "Cookie": "sid=91da0b23-7a03-4ea4-a2c7-ff1d91ada8f2; JSESSIONID=478CADC000E3BF00318FCFF2EAA5BA2F",
}
req_header_infolist = {
    # "Cookie": "sid=91da0b23-7a03-4ea4-a2c7-ff1d91ada8f2; JSESSIONID=70C893F07F0CF205848EBE81F24D4A7F",
}
req_header_scorelist = {
    # "Cookie": "sid=36393813-ff9c-4afc-bf7a-7e66885af394; JSESSIONID=CE2020478892D3C45E318861C2290C76",
}

user_agent_list = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/85.0.4183.121 Safari/537.36",
    "Mozilla/5.0 (Windows NT 6.1; rv:2.0.1) Gecko/20100101 Firefox/4.0.1",
    "Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; TencentTraveler 4.0)",
    "Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; Trident/4.0; SE 2.X MetaSr 1.0; "
    "SE 2.X MetaSr 1.0; .NET CLR 2.0.50727; SE 2.X MetaSr 1.0)",
    "Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; 360SE)",
    "Opera/9.80 (Windows NT 6.1; U; en) Presto/2.8.131 Version/11.11",
]


pro = {
    "https": "http://localhost:5000"
}


class MysqlDb(object):
    # 对存储config配置表和log配置表的操作
    def __init__(self, database_name):
        self.mysql_db = pymysql.connect(**database_name)
        self.mysql_cursor = self.mysql_db.cursor()
        self.insert_data_list = []

    def execute_sql(self, sql):
        try:
            self.mysql_cursor.execute(sql)
        except Exception as e:
            print(f'执行sql出错, {e}')
        else:
            return self.mysql_cursor.fetchone()

    def insert_data_mysql(self, table, item):
        keys, vals = list(item.keys()), list(item.values())
        field = ','.join(keys)
        val = ','.join(['%s'] * len(item))
        insert_sql = "replace into " + table + " (" + field + ") VALUES (" + val + ")"
        self.insert_data_list.append(vals)
        if len(self.insert_data_list) == 1:
            try:
                self.mysql_cursor.executemany(insert_sql, self.insert_data_list)
            except Exception as e:
                print(f'插入数据出错, {e},{item}')
            else:
                self.mysql_db.commit()
                self.insert_data_list.clear()


def partition(num_list, size):
    """
    Returns a new list with elements
    of which is a list of certain size.

        >>> partition([1, 2, 3, 4], 3)
        [[1, 2, 3], [4]]
    """
    return [num_list[i:i + size] for i in range(0, len(num_list), size)]


def thread_process(nums, url):
    pool = multiprocessing.Pool(15)
    for num in nums:
        pool.apply_async(page_id, (num, url))
    pool.close()
    pool.join()


def page_id(num, url):
    try:
        req_header_vlist['User-Agent'] = random.choice(user_agent_list)
        data_json_num = {'query': (None, '{"page": %s,"page_count": 10}' % num)}
        web_data_num = requests.post(url=url, headers=req_header_vlist, files=data_json_num)
        json_data_num = json.loads(web_data_num.text)
        # file_main_db = MysqlDb(config.MYSQL_CONFIG)
        # max_time = file_main_db.execute_sql("select MAX(pub_time) from intf_vulnerability_info ")
        # file_main_db.mysql_db.close()
        for data in json_data_num['intgs']:
            # if datetime.datetime.strptime(data['pub_time'], "%Y-%m-%d %H:%M:%S") >= max_time[0]:
            time.sleep(random.uniform(2, 4))
            next_get_detail(data['id'])
            # os._exit()
    except Exception as e:
        print(f'page_id错误信息为{num}  {e}')
        time.sleep(random.uniform(120, 300))
        page_id(num, url)


def next_get_detail(identity):
    url = 'https://redqueen.tj-un.com//Json/intelDetailsInfo.json'
    url_scorce = 'https://redqueen.tj-un.com//Json/intelDetailsHoleScoring.json'
    data_json = {'query': (None, '{"id":"%s"}' % identity)}
    req_header_infolist['User-Agent'] = random.choice(user_agent_list)
    req_header_scorelist['User-Agent'] = random.choice(user_agent_list)
    web_data = requests.post(url=url, headers=req_header_infolist, files=data_json)
    time.sleep(random.uniform(2, 5.2))
    web_data_scorce = requests.post(url=url_scorce, headers=req_header_scorelist, files=data_json)
    json_data = json.loads(web_data.text)
    json_data_scorce = json.loads(web_data_scorce.text)['hole_scoring']
    ref_info = []
    for data in json_data['ref_info']:
        ref_info.append(' '.join(list(data.values())))
    item = {
        'cve_no': json_data['title'][-14:].strip(),
        'cnnvd_name': json_data['title'],
        'pub_time': json_data['pub_time'] if json_data['pub_time'] else None,
        'pub_source': json_data['pub_source'] if json_data['pub_source'] else None,
        'ref_info': ' '.join(ref_info) if json_data['ref_info'] else None,
        'summary': ' '.join(json_data['summary']) if json_data['summary'] else None,
        'leak_type': ' '.join(json_data['hole_type']) if json_data['hole_type'] else None,
        'attacktarget': ' '.join(json_data['target']),
        'solutions': json_data['response'] if json_data['response'] else None,
        'cvss_v3_base_score': json_data_scorce[0]['val'],
        'vector': json_data_scorce[1]['val'],
        'impact_score': json_data_scorce[2]['val'],
        'exploitability_score': json_data_scorce[3]['val'],
        'attack_vector': json_data_scorce[4]['val'],
        'attack_complexity': json_data_scorce[5]['val'],
        'privileges_required': json_data_scorce[6]['val'],
        'user_interaction': json_data_scorce[7]['val'],
        'scope': json_data_scorce[8]['val'],
        'confidentiality': json_data_scorce[9]['val'],
        'integrity': json_data_scorce[10]['val'],
        'availability': json_data_scorce[11]['val'],
    }
    download_data_mysql(item)


def download_data_mysql(data):
    file_main_db = MysqlDb(config.MYSQL_CONFIG)
    isNull = file_main_db.execute_sql(f"select 1 from intf_vulnerability_info WHERE cve_no = '{data['cve_no']}'")
    if not isNull:
        if data['cve_no']:
            print(f"正在处理{data['cve_no']}的数据")
            file_main_db.insert_data_mysql('intf_vulnerability_info', data)
        else:
            print(data)
    file_main_db.mysql_db.close()
    time.sleep(random.uniform(7, 20))


def main():
    req_header_vlist['User-Agent'] = random.choice(user_agent_list)
    url = 'https://redqueen.tj-un.com/Json/intelHomeVulnIntelList.json'
    data_json = {'query': (None, '{"page": 1,"page_count": 10}')}
    web_data = requests.post(url=url, headers=req_header_vlist, files=data_json)
    time.sleep(random.uniform(2, 5))
    json_data = json.loads(web_data.text)
    n = 1  # 控制开启进程数量
    total_page = int(json_data['total']/10 + 2)
    ls = list(range(1020, total_page))
    res = partition(ls, int(floor(len(ls)/n)))
    for nums in res:
        p = multiprocessing.Process(target=thread_process, args=(nums, url, ))
        print('第{}个进程'.format(res.index(nums)))
        p.start()


if __name__ == '__main__':
    delay = 86400  # 一个天
    s.enter(delay, 2, main())
    s.run()
